package encrypt

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
)

// AES加密函数
func AESEncrypt(data, secret []byte) ([]byte, error) {
	// 创建AES块
	block, err := aes.NewCipher(secret)
	if err != nil {
		return nil, err
	}

	// 创建CBC模式
	blockSize := block.BlockSize()
	ciphertext := make([]byte, blockSize+len(data))
	iv := ciphertext[:blockSize]

	// 生成随机IV
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}

	// 填充数据
	paddedData := pkcs7Padding(data, blockSize)

	// 加密
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[blockSize:], paddedData)

	// 返回Base64编码结果
	return []byte(base64.StdEncoding.EncodeToString(ciphertext)), nil
}

// AES解密函数
func AESDecrypt(data, secret []byte) ([]byte, error) {
	// 解码Base64
	ciphertext, err := base64.StdEncoding.DecodeString(string(data))
	if err != nil {
		return nil, err
	}

	// 创建AES块
	block, err := aes.NewCipher(secret)
	if err != nil {
		return nil, err
	}

	// 检查块大小
	blockSize := block.BlockSize()
	if len(ciphertext) < blockSize {
		return nil, fmt.Errorf("密文长度不足")
	}

	// 提取IV和密文
	iv := ciphertext[:blockSize]
	ciphertext = ciphertext[blockSize:]

	// 检查密文长度
	if len(ciphertext)%blockSize != 0 {
		return nil, fmt.Errorf("密文长度不是块大小的整数倍")
	}

	// 解密
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(ciphertext, ciphertext)

	// 去除填充
	plaintext, err := pkcs7Unpadding(ciphertext)
	if err != nil {
		return nil, err
	}

	return plaintext, nil
}

// PKCS7填充
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padtext...)
}

// PKCS7去填充
func pkcs7Unpadding(data []byte) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, fmt.Errorf("数据为空")
	}

	// 获取填充字节数
	padding := int(data[length-1])
	if padding > length {
		return nil, fmt.Errorf("填充错误")
	}

	// 检查填充是否正确
	for i := 0; i < padding; i++ {
		if data[length-1-i] != byte(padding) {
			return nil, fmt.Errorf("填充错误")
		}
	}

	return data[:length-padding], nil
}
